#include "buzzlock.h"
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/unistd.h>
#include <linux/vmalloc.h>
#include <linux/sched.h>
#include <linux/jiffies.h>
long init_bzz(bzz_t *lockptr, int num_threads, useconds_t timeout) {
    //Initialize all of the structures for the lock
    printk("<1>init_bzz worked!\n");
    *lockptr = vmalloc(sizeof(bzz));
    bzz_t lock = *lockptr;
    //Initialize the locking objects[.
    sema_init(&lock -> lock, num_threads);
    sema_init(&lock -> queue_lock, 1);
    sema_init(&lock -> black_lock, 0);
    sema_init(&lock -> old_gold_lock, 0);
    init_completion(&lock -> new_gold_lock);
    //The timeout before new gold becomes old
    lock -> timeout = timeout;
    //How many threads currently have the lock
    lock -> in_use = 0;
    lock -> cur_threads = 0;
    lock -> num_threads = num_threads;
    //Counts for the number of threads of a color in the queue
    lock -> black = 0;
    lock -> new_gold = 0;
    lock -> old_gold = 0;
    //Count for the number of signals made
    lock -> bsig = 0;
    lock -> ngsig = 0;
    lock -> ogsig = 0;
    //Arrays to hold colors and threads
    lock -> colors = vmalloc(sizeof(int)*1024);
    lock -> threads = vmalloc(sizeof(int)*1024);
    
    return 0;
}

long bzz_color(int color, bzz_t lock) {
    //Color the calling thread the given color
    int i;
    down(&lock -> queue_lock);
    //Find the spot for the thread, and give it the correct color
    for(i = 0; i < 1024; i++) {
        if(lock -> threads[i] == 0) {
            lock -> threads[i] = current -> pid;
            lock -> colors[i] = color;
            break;
        } else if(lock -> threads[i] == current -> pid) {
            lock -> colors[i] = color;
            break;
        }
    }
    lock -> cur_threads++;
    DPRINT(("Thread %d colored %d\n", current->pid, color));
    up(&lock -> queue_lock);
    return 0;
}

long bzz_lock(bzz_t lock) {
    down(&lock -> queue_lock);
    int color = get_color(lock);
    DPRINT(("Found that thread %d is color %d\n", current -> pid, color));
    if(color == BZZ_BLACK) {
        lock -> black++;
        if(lock -> black == 1 && lock -> old_gold == 0 && lock -> new_gold == 0 && lock -> in_use < lock -> num_threads) {
            
            //No contest, give the lock
            lock -> in_use++; 
            lock -> black--;
            up(&lock -> queue_lock);
            DPRINT(("Aquiring lock\n"));
            down(&lock -> lock);
            DPRINT(("Black given lock.\n"));
            
        } else {
            
            up(&lock -> queue_lock);  //Release lock to allow others to queue
            DPRINT(("Black sleeping\n"));
            down(&lock -> black_lock); //Wait on black cond
            DPRINT(("Black woke up\n"));
            down(&lock -> queue_lock); //Reaquire queue lock to update data
            DPRINT(("Updated b\n"));
            lock -> in_use++;
            lock -> black--;
            lock -> bsig--;
            up(&lock -> queue_lock);
            DPRINT(("Aquiring lock\n"));
            down(&lock -> lock); //Grab the actual lock
            DPRINT(("Black aquired lock.\n"));
            
        }
    } else if(color == BZZ_GOLD) {
        lock -> new_gold++;
        if(lock -> black == 0 && lock -> old_gold == 0 && lock -> new_gold == 1 && lock -> in_use < lock -> num_threads) {
            
            //No contest, give the lock
            lock -> in_use++;
            lock -> new_gold--;
            up(&lock -> queue_lock);
            down(&lock -> lock);
            DPRINT(("New gold given lock.\n"));
            
        } else {
            up(&lock -> queue_lock);
            //If the following if returns true, we were signalled and should give the lock
            DPRINT(("New gold sleeping\n"));
            if(wait_for_completion_timeout(&lock -> new_gold_lock, usecs_to_jiffies(lock -> timeout))) {
                DPRINT(("New gold signaled\n"));
                down(&lock -> queue_lock);
                DPRINT(("Updated ng\n"));
                lock -> in_use++;
                lock -> new_gold--;
                lock -> ngsig--;
                up(&lock -> queue_lock);
                DPRINT(("Aquiring lock\n"));
                down(&lock -> lock);  //Grab the lock
                DPRINT(("New gold aquired lock.\n"));
                
            } else {
                
                DPRINT(("Timed out, making new gold old.\n"));
                down(&lock -> queue_lock); //Update old and new gold thread counts
                lock -> new_gold--;
                //If we missed a signal, we'd better grab it!
                if(lock -> new_gold == 0 && lock -> ngsig > 0) {
                    
                    lock -> ngsig--;
                    lock -> in_use++;
                    up(&lock -> queue_lock);
                    down(&lock -> lock);
                    
                } else {
                    
                    lock -> old_gold++;
                    up(&lock -> queue_lock);
                    DPRINT(("Updated og\n"));
                    down(&lock -> old_gold_lock); //Wait on old gold cond
                    DPRINT(("Old gold woke up\n"));
                    down(&lock -> queue_lock);
                    lock -> in_use++;
                    lock -> old_gold--;
                    lock -> ogsig--;
                    up(&lock -> queue_lock);
                    DPRINT(("Aquiring lock\n"));
                    down(&lock -> lock); //Grab the actual lock
                    DPRINT(("Old gold aquired lock.\n"));
                    
                }
            }
        }
    } else {
        up(&lock -> queue_lock);
    }
    return 0;
}

long bzz_release(bzz_t lock) {
    //Release the lock, and signal a waiting thread of the correct priority if needed
    DPRINT(("Releasing lock\n"));
    up(&lock -> lock);
    down(&lock -> queue_lock);
    DPRINT(("b: %d, og: %d, ng: %d\nbsig: %d, ogsig: %d, ngsig: %d\nIn use: %d\n", lock -> black, lock -> old_gold, lock -> new_gold, lock -> bsig, lock -> ogsig, lock -> ngsig, lock -> in_use));
    if(lock -> old_gold > lock -> ogsig) {
        up(&lock -> old_gold_lock);
        lock -> ogsig++;
        DPRINT(("ogsig\n"));
    } else if(lock -> black > lock -> bsig) {
        up(&lock -> black_lock);
        lock -> bsig++;
        DPRINT(("bsig\n"));
    } else if(lock -> new_gold > lock -> ngsig) {
        complete(&lock -> new_gold_lock);
        lock -> ngsig++;
        DPRINT(("ngsig\n"));
    }
    lock -> in_use--;
    up(&lock -> queue_lock);
    return 0;
}
long bzz_kill(bzz_t lock) {
    //Free malloc'd data
    DPRINT(("Killing lock\n"));
    vfree(lock -> colors);
    vfree(lock -> threads);
    vfree(lock);
    return 0;
}

int get_color(bzz_t lock) {
    //Get the color of the calling thread
    int i;
    int tid = current -> pid;
    for(i = 0; i < 1024; i++) {
        if(lock -> threads[i] == tid)
            return lock -> colors[i];
    }
    return 0;
}

int init_module() {
    //Set up the correct syscalls.
    myptrs.f1 = &init_bzz;
    myptrs.f2 = &bzz_color;
    myptrs.f3 = &bzz_lock;
    myptrs.f4 = &bzz_release;
    myptrs.f5 = &bzz_kill;
   return 0;
}

void cleanup_module() {
    //Null out pointers again
    myptrs.f1 = 0;
    myptrs.f2 = 0;
    myptrs.f3 = 0;
    myptrs.f4 = 0;
    myptrs.f5 = 0;
}

